#include "4-bit_sort_tree.hpp"

// 调整底层失衡节点，只能用于调整平衡因子相差2的树
template<class T = int>
BitNode<T>* adjustAVL(BitNode<T>* tree)
{
  if (tree == nullptr) {
    return tree;
  }

  int bal = getBalanceTr(tree);
  if (1 >= bal && bal >= -1) {
    return tree;
  }
  
  if (2 == bal) {
    int lbal = getBalanceTr(tree->lchild);
    if (-1 == lbal) {
      tree->lchild = leftRotate(tree->lchild);
    }
    return rightRotate(tree);
  }

  if (-2 == bal) {
    int rbal = getBalanceTr(tree->rchild);
    if (1 == rbal) {
      tree->rchild = rightRotate(tree->rchild);
    }
    return leftRotate(tree);
  }

  // 平衡因子超过2，不进行任何处理
  return tree;
}

// 若平衡因子超过2的树，使用后序遍历从底层开始调整
template<class T = int>
BitNode<T>* treeToAVL(BitNode<T>* tree)
{
  if (tree == nullptr) {
    return tree;
  }

  tree->lchild = treeToAVL(tree->lchild);
  tree->rchild = treeToAVL(tree->rchild);
  return adjustAVL(tree);
}

template<class T = int>
BitNode<T>* insertNodeAVL(BitNode<T>* tree, T data)
{
  if (tree == nullptr) {
    return new BitNode<T>(data);
  }

  // 添加节点不存在函数返回空指针的情况
  if (data < tree->data) {
    tree->lchild = insertNodeAVL(tree->lchild, data);
    tree->lchild->parent = tree;
  }
  if (data > tree->data) {
    tree->rchild = insertNodeAVL(tree->rchild, data);
    tree->rchild->parent = tree;
  }
  return adjustAVL(tree);
}

template<class T = int>
BitNode<T>* delNodeAVL(BitNode<T>* tree, int data)
{
  if (tree == nullptr) {
    return nullptr;
  }

  if (data == tree->data) {
    return delNodeSTr(tree);
  }

  if (data < tree->data) {
    tree->lchild = delNodeAVL(tree->lchild, data);
  }
  if (data > tree->data) {
    tree->rchild = delNodeAVL(tree->rchild, data);
  }
  return adjustAVL(tree);
}

template<class T = int>
class BalanceBitTree : public BitSortTree<T> {
  public:
    BalanceBitTree() = default;
    explicit BalanceBitTree(BitNode<T>* _tree) : BitSortTree<T>(_tree) {}
    virtual ~BalanceBitTree() = default;

    int getBalance() { return getBalanceTr(BitSortTree<T>::tree); }
    void adjust() { BitSortTree<T>::tree = adjustAVL(BitSortTree<T>::tree); }

    void insertNode(T data) override { BitSortTree<T>::tree = insertNodeAVL(BitSortTree<T>::tree, data); }
    void delNode(T data) override { BitSortTree<T>::tree = delNodeAVL(BitSortTree<T>::tree, data); }
};

